############################################################
# Part I: Setup problem with Dolfin

from dolfin import *

# Defining linear algebra package to be used (Petsc, uBlas, etc.)
parameters.linear_algebra_backend = "uBLAS"

# Read mesh from file and create function space
mesh = UnitSquare(4, 4)
V = FunctionSpace(mesh, "CG", 1)

# Define Dirichlet boundary (x = 0 or x = 1 or y = 0 or y = 1)
def boundary(x):
    return x[0] < DOLFIN_EPS or x[0] > 1.0 - DOLFIN_EPS or \
           x[1] < DOLFIN_EPS or x[1] > 1.0 - DOLFIN_EPS

# Define boundary condition
u0 = Expression('exp(x[0]*x[1])')
bc = DirichletBC(V, u0, boundary)

# Define conductivity matrix
C = as_matrix(((2.0, 1.0), (1.0, 2.0)))

# Define variational problem
v = TestFunction(V)
u = TrialFunction(V)
f = Expression("- 2*(1 + pow(x[0],2) + x[0]*x[1] + pow(x[1],2))*exp(x[0]*x[1])")
a = inner(grad(v), C*grad(u))*dx
L = v*f*dx

# Assemble matrices and vectors
A, rhs = assemble_system(a, L, bc)
############################################################



############################################################
# Part II: Solve with PyAMG
from scipy.sparse import csr_matrix
from pyamg import smoothed_aggregation_solver
(row,col,data) = A.data()   # get sparse data
n = A.size(0)
Asp = csr_matrix( (data,col.view(),row.view()), shape=(n,n))
b = rhs.data()

ml = smoothed_aggregation_solver(Asp,max_coarse=10)
residuals = []
x = ml.solve(b,tol=1e-10,accel='cg',residuals=residuals)

residuals = residuals/residuals[0]
print ml
print residuals
############################################################



############################################################
# Part III: plot
# Save solution in VTK format
#file = File("poisson.pvd")
#file << x

# Calculate max error norm
#u0_interp = interpolate(u0, V)
#u_exact = u.vector().array()
#u_numerical = u0_interp.vector().array()
#diff = abs(u_numerical - u_exact)
#print 'Max Error: %1.2E' % diff.max()

# Plot solution
#plot(x, interactive=True)

import pylab
pylab.figure(2)
pylab.semilogy(residuals)
pylab.show()
############################################################
